/* fat.c
 *
 * Copyright (C) 2010 Ricardo R. Massaro
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <dirent.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#include "fat.h"
#include "fat_dir.h"
#include "pack.h"

#define FAT_VOLUME_LABEL "WiFiDisk"

#define FAT_BPB_JMP "\xeb\x58\x90"
#define FAT_BPB_OEM "FATTER01"

#define FAT_BPB_SECTORS_PER_TRACK  63
#define FAT_BPB_HEADS              255

#define FAT_FSI_SIG1 0x41615252
#define FAT_FSI_SIG2 0x61417272
#define FAT_FSI_SIG3 0xaa550000

struct BUILD_INFO {
  FAT_INFO *fi;
  int next_free_dir_cluster;
  unsigned char *cluster_data;
};

struct ADD_DIR_INFO {
  struct BUILD_INFO *bi;
  FAT_DIR *fd;
  const char *base_dir;
};

static int add_dir(struct BUILD_INFO *bi, const char *path, int parent_cluster);

static void make_bpb(FAT_INFO *fi, unsigned char *bpb)
{
  memset(bpb, 0, 512);
  memcpy(bpb, FAT_BPB_JMP, 3);
  memcpy(bpb + 3, FAT_BPB_OEM, 8);

  pack_u16_le(bpb + 11, 512);                 // bytes per sector
  pack_u8_le(bpb + 13, 64);                   // sectors per cluster
  pack_u16_le(bpb + 14, 32);                  // reserved sectors
  pack_u8_le(bpb + 16, 2);                    // num FATs
  pack_u16_le(bpb + 17, 0);                   // num entries in root dir
  pack_u16_le(bpb + 19, 0);                   // OLD total sectors (not used)
  pack_u8_le(bpb + 21, 0xf8);                 // media type
  pack_u16_le(bpb + 22, 0);                   // OLD sectors per FAT (not used)
  pack_u16_le(bpb + 24, FAT_BPB_SECTORS_PER_TRACK);   // sectors per track
  pack_u16_le(bpb + 26, FAT_BPB_HEADS);               // heads
  pack_u32_le(bpb + 28, 0);                   // hidden sectors
  pack_u32_le(bpb + 32, fi->num_sectors);     // total sectors
  pack_u32_le(bpb + 36, fi->sectors_per_fat); // sectors per FAT
  pack_u16_le(bpb + 40, 0);                   // ext flags
  pack_u16_le(bpb + 42, 0);                   // version
  pack_u32_le(bpb + 44, 2);                   // 1st cluster of root dir
  pack_u16_le(bpb + 48, 1);                   // sector of FS info
  pack_u16_le(bpb + 50, 6);                   // sector of BPB backup
  // ... reserved ...
  pack_u8_le(bpb + 64, 0x80);                 // drive number (?)
  pack_u8_le(bpb + 65, 0);                    // reserved
  pack_u8_le(bpb + 66, 0x29);                 // boot signature
  pack_u32_le(bpb + 67, fi->volume_id);       // volume ID
  memcpy(bpb + 71, FAT_VOLUME_LABEL, 11);     // Volume Label (must also be in root dir)
  memcpy(bpb + 82, "FAT32   ", 8);            // FS type
  bpb[510] = 0x55;
  bpb[511] = 0xaa;
}

static void make_fsi(FAT_INFO *fi, unsigned char *fsi)
{
  memset(fsi, 0, 512);
  pack_u32_le(fsi,       FAT_FSI_SIG1);
  pack_u32_le(fsi + 484, FAT_FSI_SIG2);
  pack_u32_le(fsi + 488, fi->num_free_clusters);
  pack_u32_le(fsi + 492, fi->next_free_cluster);
  pack_u32_le(fsi + 508, FAT_FSI_SIG3);
}

static int make_fat(FAT_INFO *fi)
{
  fi->next_free_cluster = 2;
  fi->num_free_clusters = fi->num_clusters;

  fi->fat = malloc(512 * fi->sectors_per_fat);
  if (! fi->fat)
    return -1;
  memset(fi->fat, 0, 512 * fi->sectors_per_fat);
  pack_u32_le(fi->fat    , 0x0ffffff8);
  pack_u32_le(fi->fat + 4, 0xffffffff);
  return 0;
}

/**
 * Initialize FAT for the given disk size.
 *
 * Return:
 *  -2  invalid disk size
 *  -1  out of memory for FAT
 *   0  success
 *   1  success, but disk size too small to conform to MS spec
 */
int fat_init(FAT_INFO *fi, unsigned long long disk_size, fat_func_error error_func, fat_func_add_file add_file, void *user_data)
{
  int num_sectors, num_clusters;

  if (disk_size > 512ull * 0xffffffff)  // 2TB max since we use 512-byte sectors
    return -2;

  num_sectors = disk_size / 512;
  num_sectors -= num_sectors % (63 * 255);
  if (num_sectors < 34)  // not enough space for reserved + FATs + clusters
    return -2;

  // calculate number of clusters
  num_clusters = (num_sectors - 32) / 64;
  while (42) {    // decrease # of clusters until there's enough space for FATs
    unsigned long long req_sectors = 32 + 2 * ((4ull * (num_clusters + 2) + 511) / 512) + 64ull*num_clusters;
    if (req_sectors <= num_sectors)
      break;
    num_clusters--;
  }
  if (num_clusters == 0)
    return -2;

  // initialize
  memset(fi, 0, sizeof(FAT_INFO));
  fi->func_error = error_func;
  fi->func_add_file = add_file;
  fi->user_data = user_data;
  fi->fat = NULL;
  fi->volume_id = (unsigned int) time(NULL);
  fi->num_sectors = num_sectors;
  fi->num_clusters = num_clusters;
  fi->sectors_per_fat = (4ull * (fi->num_clusters + 2) + 511) / 512;
  fi->wasted_sectors = fi->num_sectors - 32 - 2*fi->sectors_per_fat - 64*fi->num_clusters;

  // allocate and initialize FAT
  if (make_fat(fi) != 0) {
    fi->func_error("Out of memory for FAT", NULL, fi->user_data);
    return -1;
  }

  // make BPB
  make_bpb(fi, fi->bpb);  

  if (fi->num_clusters <= 65541)  // as per MS recommendation
    return 1;
  return 0;
}

/**
 * Free memory allocated by fat_init().
 */
void fat_free(FAT_INFO *fi)
{
  if (fi->fat) {
    free(fi->fat);
    fi->fat = NULL;
  }
}

static int traverse_dir(const char *path, int (*func)(const char *path, void *p), void *p)
{
  DIR *dir;
  char buf[2048];
  struct dirent *de, *entry;

  de = (struct dirent *) buf;
  dir = opendir(path);
  if (! dir)
    return -1;
  while (readdir_r(dir, de, &entry) == 0 && entry != NULL) {
    int ret = func(entry->d_name, p);
    if (ret) {
      closedir(dir);
      return ret;
    }
  }
  closedir(dir);
  return 0;
}

static void mark_cluster_used(FAT_INFO *fi, int cluster, int next_cluster)
{
  pack_u32_le(fi->fat + 4*cluster,  next_cluster);
}

static void mark_clusters_used(FAT_INFO *fi, int first_cluster, int n_clusters)
{
  int i;

  for (i = 0; i < n_clusters; i++)
    mark_cluster_used(fi, first_cluster + i, (i == n_clusters-1) ? 0x0fffffff : first_cluster + i + 1);
  fi->next_free_cluster += n_clusters;
  fi->num_free_clusters -= n_clusters;
}

static int calc_file_num_direntries_cb(const char *file, void *p)
{
  int *p_n = (int *) p;
  const char *name;

  if (strcmp(file, ".") == 0 || strcmp(file, "..") == 0)
    return 0;

  name = strrchr(file, '/');
  if (name)
    name++; // skip '/'
  else
    name = file;

  *p_n += 1 + (strlen(name) + 12) / 13;
  return 0;
}

static int count_dir_clusters(const char *path, int is_root)
{
  int n_entries = (is_root) ? 0 : 2;    // '.' and '..' for non-root directories

  if (traverse_dir(path, calc_file_num_direntries_cb, &n_entries))
    return 1;
  return (n_entries*32 + (64*512-1)) / (64*512);
}

static int add_dir_children_cb(const char *file, void *p)
{
  struct ADD_DIR_INFO *info = (struct ADD_DIR_INFO *) p;
  struct stat st;
  char file_path[1024];

  if (strcmp(file, ".") == 0 || strcmp(file, "..") == 0)
    return 0;

  snprintf(file_path, sizeof(file_path), "%s/%s", info->base_dir, file);
  if (stat(file_path, &st) != 0) {
    info->bi->fi->func_error("Can't read information on file", file_path, info->bi->fi->user_data);
    return -1;
  }

  if ((st.st_mode & S_IFDIR) != 0) {
    // directory
    fat_dir_add_dir(info->fd, file_path, info->bi->next_free_dir_cluster);
    if (add_dir(info->bi, file_path, info->fd->first_cluster) < 0)
      return -1;
  } else if ((st.st_mode & S_IFREG) != 0) {
    // file
    int n_clusters, first_cluster;

    // add dir entries
    first_cluster = info->bi->fi->next_free_cluster;
    n_clusters = fat_dir_add_file(info->fd, file_path, first_cluster, info->bi->fi->num_free_clusters);
    if (n_clusters < 0) {
      switch (n_clusters) {
      case -1: info->bi->fi->func_error("Can't stat file", file_path, info->bi->fi->user_data); break;
      case -2: info->bi->fi->func_error("Out of space for file", file_path, info->bi->fi->user_data); break;
      default:
	info->bi->fi->func_error("Can't directory entry for file", file_path, info->bi->fi->user_data);
      }
      return -1;
    }
    
    // mark clusters used
    mark_clusters_used(info->bi->fi, first_cluster, n_clusters);
    
    // add file data
    if (n_clusters != 0 && info->bi->fi->func_add_file)
      if (info->bi->fi->func_add_file(file_path, first_cluster, n_clusters, info->bi->fi->user_data) != 0)
	return -1;
  }

  return 0;
}

static int add_dir(struct BUILD_INFO *bi, const char *path, int parent_cluster)
{
  FAT_DIR dir;
  struct ADD_DIR_INFO info;
  int i, first_cluster, n_clusters;

  n_clusters = count_dir_clusters(path, (parent_cluster) ? 0 : 1);
  
  // initialize dir struct and allocate dir clusters
  first_cluster = bi->next_free_dir_cluster;
  fat_dir_init(&dir, bi->cluster_data + (first_cluster-2)*64*512, first_cluster, n_clusters, parent_cluster);
  bi->next_free_dir_cluster += n_clusters;

  // mark directory clusters used
  for (i = 0; i < n_clusters; i++)
    mark_cluster_used(bi->fi, first_cluster + i, (i == n_clusters-1) ? 0x0fffffff : first_cluster + i + 1);

  // add '.' and '..'
  if (parent_cluster != 0) {
    fat_dir_add_dir(&dir, ".", first_cluster);
    fat_dir_add_dir(&dir, "..", parent_cluster);
  }

  // add children
  info.bi = bi;
  info.fd = &dir;
  info.base_dir = path;
  if (traverse_dir(path, add_dir_children_cb, &info))
    return -1;
  return n_clusters;
}

static int count_tree_clusters_cb(const char *dir, void *p)
{
  void **info;
  const char *base_dir;
  int *p_count;
  void *sub_info[2];
  char path[1024];
  struct stat st;
  
  if (strcmp(dir, ".") == 0 || strcmp(dir, "..") == 0)
    return 0;

  info = (void **) p;
  base_dir = (const char *) info[0];
  p_count = (int *) info[1];

  if (base_dir == NULL)
    strcpy(path, dir);
  else
    snprintf(path, sizeof(path), "%s/%s", base_dir, dir);

  // check if it's a directory
  if (stat(path, &st) != 0 || (st.st_mode & S_IFDIR) == 0)
    return 0;

  // calculate # of clusters used for dir
  *p_count += count_dir_clusters(path, base_dir == NULL);
  
  // add # of clusters used for subdirs
  sub_info[0] = (void *) path;
  sub_info[1] = (void *) p_count;
  return traverse_dir(path, count_tree_clusters_cb, sub_info);
}

int count_tree_clusters(const char *path)
{
  void *info[2];
  int n_clusters;

  info[0] = (void *) path;
  info[1] = (void *) &n_clusters;

  n_clusters = count_dir_clusters(path, 1);
  traverse_dir(path, count_tree_clusters_cb, info);
  return n_clusters;
}

/**
 * Create a FAT filesystem with the contents of the tree in 'src_dir'.
 *
 * Return the pointer to freshly allocated data containing the
 * clusters of all directories, and put the number of directory
 * clusters in *n_dir_clusters.
 */
unsigned char *fat_build_filesystem(FAT_INFO *fi, const char *src_dir, int *n_dir_clusters)
{
  struct BUILD_INFO bi;
  int n_clusters, first_cluster;
  unsigned char* cluster_data;

  // calculate # of clusters used for all directories
  n_clusters = count_tree_clusters(src_dir);

  // allocate clusters for all directories
  first_cluster = fi->next_free_cluster;
  fi->next_free_cluster += n_clusters;
  fi->num_free_clusters -= n_clusters;

  // allocate empty data for clusters
  cluster_data = (unsigned char *) malloc(n_clusters * 64*512);
  if (! cluster_data) {
    fi->func_error("Out of memory for directory clusters", NULL, fi->user_data);
    return NULL;
  }
  *n_dir_clusters = n_clusters;
  memset(cluster_data, 0, 64*512*n_clusters);

  // add directories and files
  bi.fi = fi;
  bi.cluster_data = cluster_data;
  bi.next_free_dir_cluster = first_cluster;
  if (add_dir(&bi, src_dir, 0) < 0) {
    free(cluster_data);
    return NULL;
  }

  // build FSI with free space information
  make_fsi(fi, fi->fsi);

  return cluster_data;
}

unsigned int fat_first_cluster_offset(FAT_INFO *fi)
{
  return (32 + 2*fi->sectors_per_fat) * 512;
}

unsigned int fat_cluster_size(FAT_INFO *fi)
{
  return 64 * 512;
}

unsigned long long fat_disk_size(FAT_INFO *fi)
{
  return 512ull * fi->num_sectors;
}

int fat_num_clusters(FAT_INFO *fi)
{
  return fi->num_clusters;
}

int fat_num_used_clusters(FAT_INFO *fi)
{
  return fi->next_free_cluster - 2;
}

void *fat_get_bpb(FAT_INFO *fi)
{
  return fi->bpb;
}

void *fat_get_fsi(FAT_INFO *fi)
{
  return fi->fsi;
}

unsigned int fat_fat_size(FAT_INFO *fi)
{
  return 512 * fi->sectors_per_fat;
}

void *fat_get_fat(FAT_INFO *fi, unsigned int *p_size)
{
  *p_size = 512 * fi->sectors_per_fat;
  return fi->fat;
}

